#pragma once

#include <sstream>
#include <numeric>

#include "macros/errors.h"
#include "tensor/tensor_descriptor.h"
#include "operators/operator_description_base.h"
#include "operators/names.h"

namespace txdnn {
namespace operators {
namespace rope {

struct OperatorDescriptionFwd : public OperatorDescriptionBase {
public:
  OperatorDescriptionFwd(const TensorDescriptor &x_desc,
                         const TensorDescriptor &cos_desc,
                         const TensorDescriptor &sin_desc,
                         const TensorDescriptor &y_desc)
    : x_desc_(x_desc), cos_desc_(cos_desc), sin_desc_(sin_desc), y_desc_(y_desc) {
    isValidDims();
  }

  bool isValidDims() const {
    if (x_desc_.getDims() != y_desc_.getDims()) {
      TXDNN_THROW("RoPEForward: x and y dimension length do not match.");
    }
    if (cos_desc_.getDims() != sin_desc_.getDims()) {
      TXDNN_THROW("RoPEForward: cos and sin dimension length do not match.");
    }
    return true;
  }

  bool isSameType() const {
    return x_desc_.getDataType() == cos_desc_.getDataType();
  }

  NetworkConfig makeNetworkConfig() const override {
    std::ostringstream ss;
    ss << "rope";
    auto print_strides = [&ss](const TensorDescriptor& desc) {
      for (const auto& s : desc.getStrides()) {
        ss << s << "x";
      }
    };

    auto y_dims = y_desc_.getDims();
    auto output_numel = std::accumulate(y_dims.begin(), y_dims.end(), static_cast<size_t>(1), std::multiplies<size_t>());
    auto dtype = x_desc_.getDataType();

    print_strides(x_desc_);
    print_strides(cos_desc_);
    print_strides(sin_desc_);
    print_strides(y_desc_);

    ss << "fwd-";
    ss << "dtype" << dtype;
    ss << "output_numel" << output_numel;

    return NetworkConfig(ss.str());
  }

  const TensorDescriptor& getXDesc() const { return x_desc_; }
  const TensorDescriptor& getCosDesc() const { return cos_desc_; }
  const TensorDescriptor& getSinDesc() const { return sin_desc_; }
  const TensorDescriptor& getYDesc() const { return y_desc_; }

private:
  TensorDescriptor x_desc_;
  TensorDescriptor y_desc_;
  TensorDescriptor cos_desc_;
  TensorDescriptor sin_desc_;

  // NetworkConfig makeFwdNetworkConfig() const;
};

} // namespace rope
} // namespace operators
} // namespace txdnn